//! Contains the parsing logic using the `chumsky` library.

use crate::ast::{BinaryOp, Expression, Literal, Program, Statement, UnaryOp, Vector};
use chumsky::prelude::*;
use chumsky::text;

// A type alias for the parser's extra error type.
type ParserExtraError<'a> = extra::Err<Simple<'a, char>>;

/// Parses a full MATLAB program.
pub fn program_parser<'a>() -> impl Parser<'a, &'a str, Program, ParserExtraError<'a>> {
    statement_parser()
        .repeated()
        .collect()
        .map(|statements| Program { statements })
        .then_ignore(end())
}

/// Parses a single MATLAB statement.
fn statement_parser<'a>() -> impl Parser<'a, &'a str, Statement, ParserExtraError<'a>> {
    let comment = just('%').then(none_of("\n").repeated());

    let statement_body = {
        // An assignment statement: `ident = expr;`
        let assignment = text::ident()
            .then_ignore(just('=').padded())
            .then(expression_parser())
            .map(|(name, expr)| Statement::Assignment {
                left: Expression::Identifier(name.to_string()),
                right: expr,
            });

        // An expression terminated by a semicolon or newline.
        let expr = expression_parser().map(Statement::Expression);

        // Order matters: try `assignment` first as it's more specific.
        choice((assignment, expr))
    };

    statement_body
        .then_ignore(one_of(";\n").repeated().at_least(1))
        .padded_by(comment.repeated().padded())
}

/// Parses any expression.
fn expression_parser<'a>() -> impl Parser<'a, &'a str, Expression, ParserExtraError<'a>> {
    // Recursive parser needed for parenthesized expressions and function arguments.
    recursive(|expression_parser| {
        // Parser for atomic expressions, the building blocks.
        let atom = {
            // Numbers (integers and floats)
            let num = text::digits(10)
                .then(just('.').then(text::digits(10)).or_not())
                .to_slice()
                .map(|s: &str| s.parse::<f64>().unwrap())
                .map(Literal::Number);

            // Strings
            let string = just('\'')
                .ignore_then(none_of('\'').repeated().collect::<String>())
                .then_ignore(just('\''))
                .map(Literal::String);

            // Identifiers (variable names, function names)
            let ident = text::ident().map(|s: &str| Expression::Identifier(s.to_string()));

            // Parenthesized expressions
            let parens = expression_parser
                .clone() // Cloning is necessary for recursion
                .delimited_by(just('(').padded(), just(')').padded());

            // A function call
            let func_call = text::ident()
                .then(
                    expression_parser.clone() // Clone for nested calls
                        .separated_by(just(',').padded())
                        .allow_trailing()
                        .collect::<Vec<_>>()
                        .delimited_by(just('(').padded(), just(')').padded()),
                )
                .map(|(name, args): (&'a str, Vec<_>)| Expression::FunctionCall {
                    name: name.to_string(),
                    args,
                });

            // Vector/matrix definitions
            // Parser for a single row of a matrix (e.g., `1 2, 3`)
            let row = expression_parser.clone()
                .separated_by(just(' ').or(just(',')).padded())
                .at_least(1)
                .collect::<Vec<Expression>>();

            // Parser for a non-empty matrix
            let nonempty_matrix = row
                .separated_by(just(';').padded())
                .collect::<Vec<Vec<Expression>>>()
                .map(Vector);

            // Parser for an empty matrix `[]`
            let empty_matrix = just("[]").to(Vector(vec![]));

            // Combine non-empty and empty matrix parsers, wrapped in brackets
            let vector = nonempty_matrix
                .or(empty_matrix)
                .map(Expression::Vector)
                .delimited_by(just('[').padded(), just(']').padded());


            // Atoms can be literals, identifiers, function calls, or parenthesized expressions.
            // We try parsing a function call first, then a vector, then the others.
            choice((
                func_call,
                vector,
                num.map(Expression::Literal),
                string.map(Expression::Literal),
                ident,
                parens,
            ))
            .padded()
        };

        // Unary operators (negation, logical not) are parsed recursively.
        let unary = recursive(|unary| {
            let op = just('-')
                .to(UnaryOp::Negate)
                .or(just('~').to(UnaryOp::Not));

            let operation = op
                .padded()
                .then(unary)
                .map(|(op, expr)| Expression::UnaryOp {
                    op,
                    expr: Box::new(expr),
                });

            operation.or(atom)
        });

        // Binary operators with precedence
        let op = |c, op_val| just(c).padded().to(op_val);

        let product = unary.clone().foldl(
            choice((op('*', BinaryOp::Mul), op('/', BinaryOp::Div)))
                .then(unary)
                .repeated(),
            |left, (op, right)| Expression::BinaryOp {
                op,
                left: Box::new(left),
                right: Box::new(right),
            },
        );

        let sum = product.clone().foldl(
            choice((op('+', BinaryOp::Add), op('-', BinaryOp::Sub)))
                .then(product)
                .repeated(),
            |left, (op, right)| Expression::BinaryOp {
                op,
                left: Box::new(left),
                right: Box::new(right),
            },
        );

        sum // The full expression parser
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ast::{Expression, Literal, Vector};

    // Helper function to parse an expression and unwrap the result.
    fn parse_expr(s: &str) -> Expression {
        expression_parser().parse(s).into_result().unwrap()
    }

    #[test]
    fn test_parse_row_vector() {
        let input = "[1 2 3]";
        let expected = Expression::Vector(Vector(vec![
            vec![
                Expression::Literal(Literal::Number(1.0)),
                Expression::Literal(Literal::Number(2.0)),
                Expression::Literal(Literal::Number(3.0)),
            ],
        ]));
        assert_eq!(parse_expr(input), expected);
    }

    #[test]
    fn test_parse_column_vector() {
        let input = "[1; 2; 3]";
        let expected = Expression::Vector(Vector(vec![
            vec![Expression::Literal(Literal::Number(1.0))],
            vec![Expression::Literal(Literal::Number(2.0))],
            vec![Expression::Literal(Literal::Number(3.0))],
        ]));
        assert_eq!(parse_expr(input), expected);
    }

    #[test]
    fn test_parse_matrix() {
        let input = "[1 2; 3 4]";
        let expected = Expression::Vector(Vector(vec![
            vec![
                Expression::Literal(Literal::Number(1.0)),
                Expression::Literal(Literal::Number(2.0)),
            ],
            vec![
                Expression::Literal(Literal::Number(3.0)),
                Expression::Literal(Literal::Number(4.0)),
            ],
        ]));
        assert_eq!(parse_expr(input), expected);
    }
    
    #[test]
    fn test_parse_empty_matrix() {
        let input = "[]";
        let expected = Expression::Vector(Vector(vec![]));
        assert_eq!(expression_parser().parse(input).into_result().unwrap(), expected);
    }

    #[test]
    fn test_parse_matrix_with_expressions() {
        let input = "[1 + 2, 3; 4 * 5, 6]";
        let expected = Expression::Vector(Vector(vec![
            vec![
                Expression::BinaryOp {
                    op: BinaryOp::Add,
                    left: Box::new(Expression::Literal(Literal::Number(1.0))),
                    right: Box::new(Expression::Literal(Literal::Number(2.0))),
                },
                Expression::Literal(Literal::Number(3.0)),
            ],
            vec![
                Expression::BinaryOp {
                    op: BinaryOp::Mul,
                    left: Box::new(Expression::Literal(Literal::Number(4.0))),
                    right: Box::new(Expression::Literal(Literal::Number(5.0))),
                },
                Expression::Literal(Literal::Number(6.0)),
            ],
        ]));
        assert_eq!(parse_expr(input), expected);
    }

    #[test]
    fn test_parse_matrix_with_func_call() {
        let input = "[1 1; zeros(2,1)]";
        let expected = Expression::Vector(Vector(vec![
            vec![
                Expression::Literal(Literal::Number(1.0)),
                Expression::Literal(Literal::Number(1.0)),
            ],
            vec![
                Expression::FunctionCall {
                    name: "zeros".to_string(),
                    args: vec![
                        Expression::Literal(Literal::Number(2.0)),
                        Expression::Literal(Literal::Number(1.0)),
                    ],
                },
            ],
        ]));
        assert_eq!(parse_expr(input), expected);
    }
}
